from tools9 import (minutes_to_time,time_to_minutes,parse_shift,safe_parse_list,
                    get_allowed_shifts,check_ab_vehicles, calculate_work_end_time,
                    calculate_work_start_time,calculate_unload_done_time,calculate_arrive_time)

def assign_early_aluminum_rod_tasks(df, vehicle_df, driver_df, assigned_drivers, assigned_vehicles):
    """
    分配上午铝棒任务的司机和车辆
    """
    # 1. 找出运输铝棒的车次
    assignment_results=[]  # 初始化排班结果
    # 先筛选铝棒
    aluminum_tasks = df[df['cargo_type'] == '铝棒'].copy()
    # 再筛选时间
    early_aluminum_tasks = aluminum_tasks[time_to_minutes(aluminum_tasks['departure_time']) < 721].copy()
    for _, task in early_aluminum_tasks.iterrows():
        result, df = assign_aluminum_rod_task(df, task, vehicle_df, driver_df, assigned_vehicles, assigned_drivers)

        if result is not None:
            assignment_results.append(result)
            assigned_vehicles.add(result['车辆ID'])
            assigned_drivers.add(result['司机姓名'])

    return assignment_results, assigned_drivers, assigned_vehicles,df

def assign_late_aluminum_rod_tasks(df, vehicle_df, driver_df, assigned_drivers, assigned_vehicles):
    """
    分配下午铝棒任务的司机和车辆
    """
    # 1. 找出运输铝棒的车次
    assignment_results=[]  # 初始化排班结果
    # 先筛选铝棒
    aluminum_tasks = df[df['cargo_type'] == '铝棒'].copy()
    # 再筛选时间
    early_aluminum_tasks = aluminum_tasks[time_to_minutes(aluminum_tasks['departure_time']) > 720].copy()
    for _, task in early_aluminum_tasks.iterrows():

        result,df = assign_aluminum_rod_task(df, task, vehicle_df, driver_df, assigned_vehicles, assigned_drivers)
        if result is not None:
            assignment_results.append(result)
            assigned_drivers.add(result['司机姓名'])
            assigned_vehicles.add(result['车辆ID'])

    return assignment_results, assigned_drivers, assigned_vehicles

def assign_aluminum_rod_task(df, task, vehicle_df, driver_df, assigned_vehicles, assigned_drivers):
    """
    安排铝棒任务的司机和车辆
    """
    departure_time = task['departure_time']
    task_time = departure_time
    task_shift = parse_shift(departure_time)
    vehicle_type = task['vehicle_type']


    # 先安排司机
    selected_driver = assign_driver_for_aluminum_rod(
        df,task_time, task_shift, driver_df, assigned_drivers,task
    )
    if selected_driver is None:
        print(f"没有合适司机用于铝棒任务: {departure_time}")
        return None

    # 计算到货时间
    arrive_time = calculate_arrive_time(cargo_type=task['cargo_type'],
                                        dep_time=task['departure_time'],
                                        is_second_special=False
                                        )
    # 计算卸货完成时间
    unload_done_time = calculate_unload_done_time(arrive_time=arrive_time,
                                                  cargo_type=task['cargo_type'],
                                                  load_time=task['load_time']
                                                  )
    # 计算收车时间
    end_time = calculate_work_end_time(cargo_type=task['cargo_type'],
                                       unload_done_time=unload_done_time,
                                       region=selected_driver['region'],
                                       is_transfer=task['is_transfer'],
                                       load_time=task['load_time']
                                       )

    # 再安排车辆
    if not task['plate_number']:
        available_vehicles = vehicle_df[
            (vehicle_df['type'] == vehicle_type) &
            (~vehicle_df['vehicle_id'].isin(assigned_vehicles)) &
            (vehicle_df['region'] == selected_driver['region'])
            ].copy()
        # 出车时间与收车时间间隔大于9小时
        available_time_vehicles=available_vehicles[(time_to_minutes(available_vehicles['end_time'])-time_to_minutes(selected_driver['start_time'])) < 960]
        if available_time_vehicles.empty:
            available_time_vehicles = available_vehicles

        # 计算AB角优先级
        available_vehicles['ab_priority'] = available_time_vehicles['ab_drivers'].apply(
            lambda x: check_ab_vehicles(x, selected_driver['driver_name'])
        )

        # 按AB角优先级排序（A角优先）
        available_vehicles = available_vehicles.sort_values(
            by=['ab_priority', 'vehicle_id'],
            ascending=[False, True]
        )
        if len(available_vehicles) == 0 :
            print(f"没有可用车辆用于铝棒任务: {departure_time}")
            print(task)
            return {
                '发车时间': departure_time,
                '装货时间': task['load_time'],
                '货物类型': '铝棒',
                '车辆ID': " ",
                '车辆类型': " ",
                '车牌号': " ",
                '司机ID': selected_driver['driver_id'],
                '司机姓名': selected_driver['driver_name'],
                '司机所属地区': selected_driver['region'],
                '是否接车': task['is_transfer'],
                '收车时间': end_time
            }, df
        selected_vehicle = available_vehicles.iloc[0]

    else:
        selected_vehicle = {}
        selected_vehicle['vehicle_id'] = " "
        selected_vehicle['type'] = task['vehicle_type']
        selected_vehicle['plate_number'] = task['plate_number']

    # 铝棒接车的第二车，车牌号保存
    mask = df['vehicle_id'] == task['vehicle_id']
    df.loc[mask, 'plate_number'] = selected_vehicle['plate_number']
    return {
        '发车时间': departure_time,
        '装货时间': task['load_time'],
        '货物类型': '铝棒',
        '车辆ID': selected_vehicle['vehicle_id'],
        '车辆类型': selected_vehicle['type'],
        '车牌号': selected_vehicle['plate_number'],
        '司机ID': selected_driver['driver_id'],
        '司机姓名': selected_driver['driver_name'],
        '司机所属地区': selected_driver['region'],
        '是否接车': task['is_transfer'],
        '收车时间': end_time
    },df

def assign_driver_for_aluminum_rod(df, task_time, task_shift, driver_df, assigned_drivers,task):
    """
    为铝棒任务安排司机 - 特殊处理常德司机限制
    """
    # 获取所有铝棒任务的时间，检查是否为最后一车
    all_rod_times = []
    for _, row in df[df['cargo_type'] == '铝棒'].iterrows():
        time_val = time_to_minutes(row['departure_time'])
        if time_val != -1:
            all_rod_times.append(time_val)

    last_rod_time = max(all_rod_times) if all_rod_times else 0
    is_last_task = (task_time == last_rod_time and task_time > minutes_to_time(721))

    # 筛选符合条件的司机
    eligible_drivers = driver_df[
        (~driver_df['driver_name'].isin(assigned_drivers))
        ].copy()

    if eligible_drivers.empty:
        return None
    # 计算允许的班次
    eligible_drivers['history_shifts'] = eligible_drivers['last_4days_departure_times'].apply(
        lambda x: [parse_shift(t) for t in safe_parse_list(x) if t is not None]
    )
    eligible_drivers['allowed_shifts'] = eligible_drivers['history_shifts'].apply(get_allowed_shifts)

    # 筛选允许当前班次的司机
    eligible_drivers = eligible_drivers[
        eligible_drivers['allowed_shifts'].apply(lambda x: task_shift in x if x else False)
    ]
    # 筛选车型与司机匹配
    if task['vehicle_type'] != 'SCANIA':
        eligible_drivers = eligible_drivers[eligible_drivers['region'] != '常德']

    # 如果是最后一车铝棒，排除常德司机
    if is_last_task:
        eligible_drivers = eligible_drivers[eligible_drivers['region'] != '常德']

    if eligible_drivers.empty:
        return None

    # 计算时间优先级
    eligible_drivers['start_time'] = eligible_drivers['region'].apply(
        lambda x:calculate_work_start_time(load_time=task['load_time'],
        departure_time = task['departure_time'],
        region=x,
        cargo_type=task['cargo_type'])
    )

    eligible_drivers['time_priority'] = time_to_minutes(eligible_drivers['start_time']) - time_to_minutes(eligible_drivers['end_time'])

    # 按时间优先级排序，然后按司机ID排序
    eligible_drivers = eligible_drivers.sort_values(
        by=['time_priority', 'driver_id'],
        ascending=[False, True]
    )

    return eligible_drivers.iloc[0]

